package com.starhub.engine.milvus.service.impl;

import com.google.common.collect.Lists;
import com.starhub.application.rag.biz.MilvusClientFactory;
import com.starhub.engine.milvus.service.IMilvusService;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.*;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.highlevel.dml.DeleteIdsParam;
import io.milvus.param.highlevel.dml.response.DeleteResponse;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Service
public class MilvusServiceImpl implements IMilvusService {

    private static final Logger log = LoggerFactory.getLogger(MilvusServiceImpl.class);

    private final ThreadLocal<String> threadLocalMark = new ThreadLocal<>();
    private final ThreadLocal<MilvusServiceClient> threadLocalClient = new ThreadLocal<>();

    @Autowired
    private MilvusClientFactory milvusClientFactory;

    final IndexType INDEX_TYPE = IndexType.IVF_FLAT;   // IndexType
    final String INDEX_PARAM = "{\"nlist\":1024}";     // ExtraParam

    /**
     * 创建集合的字段
     * text_id  对应的文本id
     * vector  向量字段
     * tag  标签
     */
    private final String TEXTID = "text_id";
    private final String VECTOR = "vector";
    private final String TAG = "tag";

    private final int dimension = 1024;

    @Override
    public void setMark(String mark) {
        MilvusServiceClient client = milvusClientFactory.getClient(mark);
        threadLocalClient.set(client);
        threadLocalMark.set(mark);
    }

    private String getMark() {
        String mark = threadLocalMark.get();
        if (mark == null || mark.isEmpty()) {
            throw new IllegalStateException("Mark is not set for the current thread.");
        }
        return mark;
    }

    private MilvusServiceClient getClient() {
        MilvusServiceClient client = threadLocalClient.get();
        String mark = getMark();
        if (client == null) {
            throw new IllegalStateException("No Milvus client found for mark: " + mark);
        }
        return client;
    }

    @Override
    public void create(String collectionName, String desc) {
        log.info("Miluvs create collectionName:{}, desc:{}", collectionName, desc);
        boolean has = hasCollect(collectionName);
        log.info("Miluvs hasCollect:{}", has);
        // 不存在此集合才进行创建集合
        if(!has) {
            //  创建集合 设置索引 加载集合到内存中
            FieldType fieldType1 = FieldType.newBuilder()
                    .withName(TEXTID)
                    .withDataType(DataType.Int64)
                    .withPrimaryKey(true)
                    .withAutoID(false)
                    .build();
            FieldType fieldType2 = FieldType.newBuilder()
                    .withName(VECTOR)  // 设置向量名称
                    .withDataType(DataType.FloatVector)  // 设置向量类型
                    .withDimension(dimension) // 设置向量维度
                    .build();
            FieldType fieldType3 = FieldType.newBuilder()
                    .withName(TAG)
                    .withDataType(DataType.Int64)
                    .build();
            CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withDescription(desc)
                    .withShardsNum(2)
                    .addFieldType(fieldType1)
                    .addFieldType(fieldType2)
                    .addFieldType(fieldType3)
                    .withEnableDynamicField(true)
                    .withConsistencyLevel(ConsistencyLevelEnum.BOUNDED)
                    .build();
            R<RpcStatus> response = getClient().createCollection(createCollectionReq);
            if (response.getStatus() != R.Status.Success.getCode()) {
                log.info("milvus create fail message:{}", response.getMessage());
            } else {
                // 创建集合索引并加载集合到内存
                createIndex(collectionName);
            }
        }
    }

    /**
     * 创建集合索引 -- 加在向量字段上
     * @param collectionName 集合名称
     */
    private void createIndex(String collectionName) {
        getClient().createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(collectionName)
                        .withFieldName(VECTOR)
                        .withIndexType(INDEX_TYPE)
                        .withMetricType(MetricType.L2)
                        .withExtraParam(INDEX_PARAM)
                        .withSyncMode(Boolean.FALSE)
                        .build()
        );
        // 加载所创建的集合
        loadCollection(collectionName);
    }

    /**
     * 加载集合
     * @param collectionName 集合名称
     */
    private void loadCollection(String collectionName) {
        getClient().loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .build()
        );
        // 检查加载状态
        GetLoadStateParam param = GetLoadStateParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        R<GetLoadStateResponse> stateResponse = getClient().getLoadState(param);
        if (stateResponse.getStatus() != R.Status.Success.getCode()) {
            log.error("Load collection failed: {}", stateResponse.getMessage());
        }
    }

    @Override
    public boolean hasCollect(String collectionName) {
        R<Boolean> hasResult = getClient().hasCollection(
                HasCollectionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .build());
        if (hasResult.getStatus() == R.Status.Success.getCode()) {
            return hasResult.getData();
        }
        return false;
    }

    @Override
    public boolean insert(String collectionName, List<Map<String, Object>> dataList) {
        log.info("milvus insert collectionName:{}, dataList:{}", collectionName, dataList);
        List<Long> textIds = new ArrayList<>();
        List<List<Float>> vectorList = new ArrayList<>();
        List<Long> tagList = new ArrayList<>();

        for (Map<String, Object> data : dataList) {
            textIds.add(Long.valueOf(data.get(TEXTID).toString()));
            vectorList.add((List<Float>) data.get(VECTOR));
            tagList.add(0L); // 默认标签为0
        }

        List<InsertParam.Field> fieldsInsert = new ArrayList<>();
        fieldsInsert.add(new InsertParam.Field(TEXTID, textIds));
        fieldsInsert.add(new InsertParam.Field(VECTOR, vectorList));
        fieldsInsert.add(new InsertParam.Field(TAG, tagList));

        InsertParam param = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fieldsInsert)
                .build();

        R<MutationResult> response = getClient().insert(param);
        return response.getStatus() == R.Status.Success.getCode();
    }

    @Override
    public List<Map<String, Object>> search(String collectionName, String vector, Integer topK) {
        // 将字符串形式的向量转换为List<Float>
        List<Float> vectorData = parseVectorString(vector);
        List<List<Float>> vectorList = new ArrayList<>();
        vectorList.add(vectorData);

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withVectorFieldName(VECTOR)
                .withOutFields(Lists.newArrayList("*"))
                .withVectors(vectorList)
                .withTopK(topK)
                .build();

        R<SearchResults> searchResults = getClient().search(searchParam);
        if (searchResults.getStatus() != R.Status.Success.getCode()) {
            log.error("Search failed: {}", searchResults.getMessage());
            return new ArrayList<>();
        }

        List<Map<String, Object>> results = new ArrayList<>();
        SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResults.getData().getResults());
        for (int i = 0; i < vectorList.size(); ++i) {
            List<SearchResultsWrapper.IDScore> scores = wrapper.getIDScore(i);
            for (SearchResultsWrapper.IDScore score : scores) {
                results.add(score.getFieldValues());
            }
        }
        return results;
    }

    @Override
    public boolean delete(String collectionName, List<String> ids) {
        String expr = TEXTID + " in " + ids;
        DeleteParam param = DeleteParam.newBuilder()
                .withCollectionName(collectionName)
                .withExpr(expr)
                .build();
        R<MutationResult> response = getClient().delete(param);
        return response.getStatus() == R.Status.Success.getCode();
    }

    @Override
    public boolean update(String collectionName, Map<String, Object> updateData) {
        // Milvus目前不支持直接更新向量数据，需要先删除后插入
        String id = updateData.get(TEXTID).toString();
        boolean deleteResult = delete(collectionName, Lists.newArrayList(id));
        if (!deleteResult) {
            return false;
        }
        return insert(collectionName, Lists.newArrayList(updateData));
    }

    @Override
    public List<Map<String, Object>> get(String collectionName, List<String> ids) {
        String expr = TEXTID + " in " + ids;
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withExpr(expr)
                .withOutFields(Lists.newArrayList("*"))
                .build();

        R<SearchResults> searchResults = getClient().search(searchParam);
        if (searchResults.getStatus() != R.Status.Success.getCode()) {
            log.error("Get data failed: {}", searchResults.getMessage());
            return new ArrayList<>();
        }

        List<Map<String, Object>> results = new ArrayList<>();
        SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResults.getData().getResults());
        List<SearchResultsWrapper.IDScore> scores = wrapper.getIDScore(0);
        for (SearchResultsWrapper.IDScore score : scores) {
            results.add(score.getFieldValues());
        }
        return results;
    }

    @Override
    public Long count(String collectionName) {
        R<GetCollectionStatisticsResponse> countResponse = getClient().getCollectionStatistics(
                GetCollectionStatisticsParam.newBuilder()
                        .withCollectionName(collectionName)
                        .build());
        if (countResponse.getStatus() == R.Status.Success.getCode()) {
            return countResponse.getData().getStatsList().stream()
                    .filter(stat -> "row_count".equals(stat.getKey()))
                    .findFirst()
                    .map(stat -> Long.parseLong(stat.getValue()))
                    .orElse(0L);
        }
        return 0L;
    }

    @Override
    public boolean drop(String collectionName) {
        R<RpcStatus> response = getClient().dropCollection(
                DropCollectionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .build());
        return response.getStatus() == R.Status.Success.getCode();
    }

    @Override
    public List<String> listCollections() {
        R<ShowCollectionsResponse> response = getClient().showCollections(
                ShowCollectionsParam.newBuilder().build());
        if (response.getStatus() == R.Status.Success.getCode()) {
            return response.getData().getCollectionNamesList();
        }
        return new ArrayList<>();
    }

    /**
     * 解析字符串形式的向量数据
     * @param vectorString 向量字符串，格式如："[0.1,0.2,0.3]"
     * @return 向量数据列表
     */
    private List<Float> parseVectorString(String vectorString) {
        // 移除首尾的方括号
        String content = vectorString.trim().substring(1, vectorString.length() - 1);
        String[] values = content.split(",");
        List<Float> vector = new ArrayList<>();
        for (String value : values) {
            vector.add(Float.parseFloat(value.trim()));
        }
        return vector;
    }
}
